import torchvision
from torch import nn
from torch.utils.data import DataLoader

datasets = torchvision.datasets.CIFAR10(root='./data/data_non_linear', train=False, download=True,
                                        transform=torchvision.transforms.ToTensor())

train_loader = DataLoader(datasets, batch_size=1)


class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 5, padding=2)
        self.maxpool1 = nn.MaxPool2d(kernel_size=2)
        self.conv2 = nn.Conv2d(32, 32, 5, padding=2)
        self.maxpool2 = nn.MaxPool2d(kernel_size=2)
        self.conv3 = nn.Conv2d(32, 64, 5, padding=2)
        self.maxpool3 = nn.MaxPool2d(kernel_size=2)
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(1024, 64)
        self.linear2 = nn.Linear(64, 10)

        # self.model1 = nn.Sequential(self.conv1, self.maxpool1, self.conv2, self.maxpool2, self.conv3, self.maxpool3,
        #                             self.flatten, self.linear1, self.linear2)

    def forward(self, x):
        # x = self.conv1(x)
        x = self.maxpool1(self.conv1(x))
        # x = self.conv2(x)
        x = self.maxpool2(self.conv2(x))
        # x = self.conv3(x)
        x = self.maxpool3(self.conv3(x))
        x = self.flatten(x)
        x = self.linear1(x)
        x = self.linear2(x)
        return x


loss = nn.CrossEntropyLoss()
net = SimpleNN()
for data in train_loader:
    inputs, labels = data
    outputs = net(inputs)
    # print(f"输出output：{outputs}")
    # print(f"标签labels：{labels}")
    loss_value = loss(outputs, labels)
    # print(f"损失loss_value：{loss_value}")
    loss_value.backward()
    print("ok")
